package com.ztom.top100;

import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * 数组中第 K 个最大的元素
 * <p>
 * https://leetcode-cn.com/problems/kth-largest-element-in-an-array/
 *
 * @author ZhangTao
 */
public class Code63FindKthLargest {

    /**
     * 小根堆
     *
     * @param nums
     * @param k
     * @return
     */
    public int findKthLargest1(int[] nums, int k) {
        if (nums == null || nums.length == 0 || k > nums.length) {
            return 0;
        }
        PriorityQueue<Integer> heap = new PriorityQueue<>(Comparator.comparingInt(a -> a));
        // 先放入前 k 个数据
        for (int i = 0; i < k; i++) {
            heap.add(nums[i]);
        }
        for (int i = k; i < nums.length; i++) {
            if (heap.peek() < nums[i]) {
                heap.poll();
                heap.add(nums[i]);
            }
        }
        return heap.peek();
    }

    /**
     * 改写快排
     *
     * @param nums
     * @param k
     * @return
     */
    public int findKthLargest(int[] nums, int k) {
        // 第 k 大对应数组下标 nums.length - k
        k = nums.length - k;
        int l = 0;
        int r = nums.length - 1;
        int[] mid = null;
        int pivot = 0;
        while (l <= r) {
            pivot = nums[l + (int) (Math.random() * (r - l + 1))];
            mid = partition(nums, l, r, pivot);
            if (k < mid[0]) {
                r = mid[0] - 1;
            } else if (k > mid[1]) {
                l = mid[1] + 1;
            } else {
                return pivot;
            }
        }
        return -1;
    }

    private int[] partition(int[] nums, int l, int r, int pivot) {
        if (l >= r) {
            return new int[]{l, r};
        }
        int less = l - 1;
        int more = r + 1;
        int index = l;
        while (index < more) {
            if (nums[index] == pivot) {
                index++;
            } else if (nums[index] < pivot) {
                swap(nums, ++less, index++);
            } else {
                swap(nums, --more, index);
            }
        }
        return new int[]{++less, --more};
    }

    private void swap(int[] nums, int a, int b) {
        int tmp = nums[a];
        nums[a] = nums[b];
        nums[b] = tmp;
    }
}
